# Robotics
''' Quantities:
         SE3: 4x4 homogeneous coordinates
         se3: 6x1 twist coordinate of SE3
         SO3: 3x3
'''
import math
import numpy as np
from pyquaternion import Quaternion
from math import cos, acos, sin, asin, pi, atan2

def qtn2SO3(qtn):
    qtn = qtnNorm(qtn)
    x, y, z, w = qtn.x, qtn.y, qtn.z, qtn.w
    rot_matrix00 = 1 - 2 * y * y - 2 * z * z
    rot_matrix01 = 2 * x * y - 2 * w * z
    rot_matrix02 = 2 * x * z + 2 * w * y
    rot_matrix10 = 2 * x * y + 2 * w * z
    rot_matrix11 = 1 - 2 * x * x - 2 * z * z
    rot_matrix12 = 2 * y * z - 2 * w * x
    rot_matrix20 = 2 * x * z - 2 * w * y
    rot_matrix21 = 2 * y * z + 2 * w * x
    rot_matrix22 = 1 - 2 * x * x - 2 * y * y
    return np.asarray([
        [rot_matrix00, rot_matrix01, rot_matrix02],
        [rot_matrix10, rot_matrix11, rot_matrix12],
        [rot_matrix20, rot_matrix21, rot_matrix22]
    ], dtype=np.float64)

def SO32qtn(SO3):
    qtn = Quaternion(matrix=SO3)
    q = np.quaternion(0)
    q.w = qtn.w
    q.x = qtn.x
    q.y = qtn.y
    q.z = qtn.z
    return qtnNorm(q)

def angBTqtn(qtn0, qtn1):
    '''
    :param qtn0: quaternion 1
    :param qtn1: quaternion 2
    :return: the minimum angle between 2 quaternions
    '''
    qtn0, qtn1 = qtnNorm(qtn0), qtnNorm(qtn1)
    dot = qtn0.w * qtn1.w + qtn0.x * qtn1.x + qtn0.y * qtn1.y + qtn0.z * qtn1.z
    cos_value = 2.0 * dot * dot - 1.0
    if cos_value > 0.99999:
        ang = 0
    elif cos_value < -0.99999:
        ang = pi
    else:
        ang = acos(cos_value)
    if ang > pi:
        ang = 2.0 * pi - ang
    return abs(ang)

def qtnSlerp(qtn0, qtn1, t):
    qtn0, qtn1 = qtnNorm(qtn0), qtnNorm(qtn1)
    dot = qtn0.w * qtn1.w + qtn0.x * qtn1.x + qtn0.y * qtn1.y + qtn0.z * qtn1.z
    absD = abs(dot)
    one = 0.99999
    if absD >= one:
        scale0 = 1 - t
        scale1 = t
    else:
        theta = acos(absD)
        sinTheta = sin(theta)
        scale0 = sin((1 - t) * theta) / sinTheta
        scale1 = sin(t * theta) / sinTheta
    if dot < 0:
        scale1 = -scale1
    qt = scale0 * qtn0 + scale1 * qtn1
    return qtnNorm(qt)

def qtnLerp(qtn0, qtn1, t):
    qtn0, qtn1 = qtnNorm(qtn0), qtnNorm(qtn1)
    qt = (1-t) * qtn0 + t * qtn1
    return qtnNorm(qt)

def Lerp(q0, q1, t):
    qt = (1 - t) * q0 + t * q1

    return qtnNorm(qt)

def qtnNorm(qtn):
    denominator = math.sqrt(pow(qtn.w, 2) + pow(qtn.x, 2) + pow(qtn.y, 2) + pow(qtn.z, 2))
    return qtn * (1/denominator)


def qtn2AngAixs(qtn):
    denominator = math.sqrt(pow(qtn.x, 2) + pow(qtn.y, 2) + pow(qtn.z, 2))


def wedge4(twist):
    return np.array([[0, -twist[0], -twist[1], -twist[2]],
                     [twist[0],  0,  twist[2], -twist[1]],
                     [twist[1],  -twist[2], 0,  twist[0]],
                     [twist[2],  twist[1], -twist[0], 0]])

def wedge(v):
    return np.array([[  0, -v[2],  v[1]],
                     [ v[2],    0, -v[0]],
                     [-v[1],  v[0],    0]])

def SE32Adj(SE3):
    Adj = np.zeros((6, 6))
    Adj[:3, :3] = SE3[:3, :3]
    Adj[3:, 3:] = SE3[:3, :3]
    Adj[:3, 3:] = wedge(SE3[:3, 3]) @ SE3[:3, :3]
    return Adj

def posem2SE3(pose):
    SE3 = np.identity(4)
    SE3[0, 3] = pose.pos[0]
    SE3[1, 3] = pose.pos[1]
    SE3[2, 3] = pose.pos[2]
    SE3[:3, :3] = qtn2SO3(pose.qtn)
    return SE3

def xyz2qtn(euler_ang):
    """
      Convert an Euler angle to a quaternion.

      Input
        :param euler_ang: The roll (rotation around x-axis) angle ; The pitch (rotation around y-axis) angle ; The yaw (rotation around z-axis) angle in radians.

      Output
        :return qx, qy, qz, qw: The orientation in quaternion [x,y,z,w] format
      """
    qx = np.sin(euler_ang[0] / 2) * np.cos(euler_ang[1] / 2) * np.cos(euler_ang[2] / 2) - np.cos(euler_ang[0] / 2) * np.sin(euler_ang[1] / 2) * np.sin(euler_ang[2] / 2)
    qy = np.cos(euler_ang[0] / 2) * np.sin(euler_ang[1] / 2) * np.cos(euler_ang[2] / 2) + np.sin(euler_ang[0] / 2) * np.cos(euler_ang[1] / 2) * np.sin(euler_ang[2] / 2)
    qz = np.cos(euler_ang[0] / 2) * np.cos(euler_ang[1] / 2) * np.sin(euler_ang[2] / 2) - np.sin(euler_ang[0] / 2) * np.sin(euler_ang[1] / 2) * np.cos(euler_ang[2] / 2)
    qw = np.cos(euler_ang[0] / 2) * np.cos(euler_ang[1] / 2) * np.cos(euler_ang[2] / 2) + np.sin(euler_ang[0] / 2) * np.sin(euler_ang[1] / 2) * np.sin(euler_ang[2] / 2)

    return np.array([qw, qx, qy, qz])

def qtn2xyz(qtn):
    x = atan2(2*(qtn.w * qtn.x + qtn.y * qtn.z), 1 - 2 * (qtn.x * qtn.x + qtn.y * qtn.y))
    y = asin(2 * (qtn.w * qtn.y - qtn.z * qtn.x))
    z = atan2(2*(qtn.w * qtn.z + qtn.x * qtn.y), 1 - 2 * (qtn.y * qtn.y + qtn.z * qtn.z))
    # x1 = math.degrees(x)
    # y1 = math.degrees(y)
    # z1 = math.degrees(z)
    return np.array([x, y, z])
class PoseRobot(object):
    def __init__(self, qtn=Quaternion(), pos=np.array([0, 0, 0])):
        self.qtn = qtn
        self.pos = pos
    def getSO3(self):
        return qtn2SO3(self.qtn)

    def PoseInit(self, pose):
        self.qtn = Quaternion(np.array(pose[3:]))
        self.pos = np.array(pose[:3])









